import random
import RL_agents.PPO as PPO
import Models.agent_models as models
import pygame
import simulation
import cv2
import json


# end_point_list = ((24, 531),(310, 485),(362, 90),(567, 331))
end_point_list = ((24, 531),(331, 421),(362, 90),(567, 331))
MODEL_CLASS =  models.Resnet_18
NAME        =  "time_capsule/MINI_RESNET18_unreset_2500"
IMG_SHAPE   =  (96, 96)
MAX_STEP    =  300
VAL_TIMES   =  100
Raduis      =  6.5


def img_resize(img, shape):
	# cv2的 x, y 坐标与 np 的相反
	img = cv2.resize(img, (shape[1], shape[0]))
	img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

	img = img.reshape(1, shape[0], shape[1])
	return 2. * img / 255 - 1

def main():
	gym = simulation.Enviroment()
	gym.bg = pygame.image.load("img_rsc\\3.png")
	gym.radius = Raduis
	# 读取血管路径
	with open("real.json") as json_file:
		gym.wall_poslist = json.load(json_file)
	# 初始化虚拟环境
	s = gym.reset()
	gym.max_steps = MAX_STEP
	
	model_name = NAME
	model = MODEL_CLASS(input_channels=1, act_num=5)
	input_shape = IMG_SHAPE
	agent = PPO.PPO_agent()
	agent.load_model(model)
	
	try:
		agent.load_weights(f"torch_weights/{model_name}.pt")
		print("Trained Model")
	except Exception as e:
		print(f"{e}: New Model")
	
	epoch     = VAL_TIMES                        # 训练次数
	begin_epo = 0
	for epo in range(begin_epo, begin_epo + epoch):

		end_point = random.choice(end_point_list)
		s = gym.reset(end_point, True)
		s = img_resize(s, input_shape)
		total = 0		
		for _ in range(gym.max_steps + 1):
			a, _ = agent.select_action(s)
			s_next, r, d, debug = gym.step(a)
			s_next = img_resize(s_next, input_shape)
			d = d or gym.beautiful_render()					
			s = s_next
			total += r
			if d:
				gym.close()
				break
		print(f"#epo{epo} step: {debug[2]}\tscore: {total}")


if __name__ == '__main__':
	main()
